"""
数据库模块
从database_core.py导入所有功能，避免与database目录的命名冲突
"""

from sqlalchemy import declarative_base
import sqlalchemy as sa
from utils.logger import logger

# 从database_core.py导入所有必要的功能
from .database_core import (
    get_db,
    async_session_factory as async_session_maker,
    db_manager,
    init_db,
    close_db,
    KnowledgeBase,
    AgentTask,
    UserInteraction,
    check_and_install_pgvector
)

# 定义__all__列表，明确导出的内容
__all__ = [
    'get_db', 
    'async_session_maker', 
    'db_manager', 
    'Base', 
    'init_db', 
    'close_db',
    'KnowledgeBase',
    'AgentTask',
    'UserInteraction',
    'check_and_install_pgvector'
]

# 创建基类
Base = declarative_base()

# 向量扩展类型
try:
    from pgvector.sqlalchemy import Vector
    VECTOR_TYPE_AVAILABLE = True
except ImportError:
    logger.warning("pgvector未安装，向量检索功能将不可用")
    VECTOR_TYPE_AVAILABLE = False
    
    # 创建一个模拟的Vector类型以避免导入错误
    class Vector(sa.types.UserDefinedType):
        def __init__(self, dimension=None):
            self.dimension = dimension
        
        def get_col_spec(self):
            return f'Vector({self.dimension})' if self.dimension else 'Vector'
        
        def bind_processor(self, dialect):
            return lambda x: x
        
        def result_processor(self, dialect, coltype):
            return lambda x: x


async def check_pgvector_support() -> bool:
    """
    检查数据库是否支持PgVector扩展
    """
    try:
        async for session in get_async_db_session():
            async with session.begin():
                # 检查vector扩展是否已安装
                result = await session.execute(
                    text("SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'vector')")
                )
                exists = result.scalar()
                
                if exists and VECTOR_TYPE_AVAILABLE:
                    # 测试vector类型是否可用
                    test_query = text("CREATE TEMP TABLE IF NOT EXISTS vector_test (embedding vector(3))")
                    await session.execute(test_query)
                    await session.execute(text("DROP TABLE IF EXISTS vector_test"))
                    return True
                
                return False
    
    except Exception as e:
        logger.error(f"检查PgVector支持失败: {str(e)}")
        return False


async def _ensure_extensions():
    """确保PgVector扩展已安装"""
    if not VECTOR_TYPE_AVAILABLE:
        logger.warning("跳过PgVector扩展检查，因为模块未安装")
        return
    
    try:
        async for session in get_async_db_session():
            async with session.begin():
                # 检查扩展是否存在
                result = await session.execute(
                    text("SELECT EXISTS(SELECT 1 FROM pg_extension WHERE extname = 'vector')")
                )
                extension_exists = result.scalar()
                
                if not extension_exists:
                    logger.info("正在创建pgvector扩展...")
                    await session.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
                    logger.info("pgvector扩展创建成功")
                else:
                    logger.info("pgvector扩展已存在")
    except Exception as e:
        logger.warning(f"检查/创建pgvector扩展时出错: {str(e)}")


# 导入备份系统
from core.backup.backup_manager import initialize_backup_manager
from core.backup.backup_scheduler import start_backup_scheduler

# 兼容层 - 保留DatabaseManager类以兼容旧代码
class DatabaseManager:
    """数据库管理器（兼容层）"""
    
    @property
    def engine(self):
        return async_db_pool_manager.engine
    
    @property
    def async_session_factory(self):
        return async_db_pool_manager.async_session_factory
    
    async def initialize(self):
        # 确保连接池已初始化
        if not self.engine:
            await async_db_pool_manager.initialize()
        # 确保扩展已安装
        await _ensure_extensions()
        logger.info("数据库连接初始化成功（使用优化的连接池）")
        return self.engine
    
    async def close(self):
        await async_db_pool_manager.close()
    
    async def _ensure_extensions(self):
        await _ensure_extensions()


# 创建全局数据库管理器实例（兼容层）
db_manager = DatabaseManager()


async def init_db():
    """初始化数据库并集成备份系统"""
    # 确保连接池已初始化
    if not async_db_pool_manager.engine:
        await async_db_pool_manager.initialize()
    
    # 确保扩展已安装
    await _ensure_extensions()
    
    # 创建所有表
    async with async_db_pool_manager.engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    
    # 初始化备份管理器
    await initialize_backup_manager()
    
    # 启动备份调度器
    await start_backup_scheduler()
    
    logger.info("数据库表初始化完成")


async def close_db():
    """关闭数据库连接"""
    await async_db_pool_manager.close()


async def get_db():
    """获取数据库会话（用于FastAPI依赖注入）"""
    async for session in get_async_db_session():
        yield session


# 数据库模型定义
class KnowledgeBase(Base):
    """知识库模型"""
    __tablename__ = "knowledge_base"
    
    id = sa.Column(sa.Integer, primary_key=True, index=True)
    content = sa.Column(sa.Text, nullable=False)
    embedding = sa.Column(Vector(1536), nullable=False)  # 假设使用OpenAI的embedding维度
    metadata_ = sa.Column(sa.JSON, nullable=True)
    created_at = sa.Column(sa.DateTime(timezone=True), server_default=sa.func.now())
    updated_at = sa.Column(sa.DateTime(timezone=True), onupdate=sa.func.now())
    
    # 创建向量索引
    __table_args__ = (
        sa.Index('idx_embedding', 'embedding', postgresql_using='ivfflat',
                postgresql_with={'lists': 100}, postgresql_ops={'embedding': 'vector_cosine'}),
    )


class AgentTask(Base):
    """智能体任务模型"""
    __tablename__ = "agent_tasks"
    
    id = sa.Column(sa.String(36), primary_key=True, index=True)  # UUID
    task_type = sa.Column(sa.String(50), nullable=False, index=True)
    status = sa.Column(sa.String(20), nullable=False, default='pending')
    input_data = sa.Column(sa.JSON, nullable=False)
    output_data = sa.Column(sa.JSON, nullable=True)
    agent_id = sa.Column(sa.String(100), nullable=False)
    priority = sa.Column(sa.Integer, default=0)
    created_at = sa.Column(sa.DateTime(timezone=True), server_default=sa.func.now())
    updated_at = sa.Column(sa.DateTime(timezone=True), onupdate=sa.func.now())
    completed_at = sa.Column(sa.DateTime(timezone=True), nullable=True)


class UserInteraction(Base):
    """用户交互模型"""
    __tablename__ = "user_interactions"
    
    id = sa.Column(sa.Integer, primary_key=True, index=True)
    user_id = sa.Column(sa.String(100), nullable=False, index=True)
    query = sa.Column(sa.Text, nullable=False)
    response = sa.Column(sa.Text, nullable=False)
    session_id = sa.Column(sa.String(100), nullable=False, index=True)
    feedback_score = sa.Column(sa.Integer, nullable=True)
    created_at = sa.Column(sa.DateTime(timezone=True), server_default=sa.func.now())
    query_embedding = sa.Column(Vector(1536), nullable=True)
    
    # 创建向量索引
    __table_args__ = (
        sa.Index('idx_query_embedding', 'query_embedding', postgresql_using='ivfflat',
                postgresql_with={'lists': 100}, postgresql_ops={'query_embedding': 'vector_cosine'}),
    )
